from Cryptodome.Util.number import getPrime, long_to_bytes
from random import randint

p = 113641842374773044745452662471378715712124744803786300890087741266845552497445274911054126455206043851172654951999459532631143246187756241338954515000426316814056710573904671705426280457773213203770699075812543476488669211671733660131140193807192700380376595760637713843721593216448687060984159575308383246829
g = 4971905734578114578103012188534162412154217839402427699027620871253633102571835127029750314112156046747872303795577123499710901658849488125954816314744943661146055519239143913435154769457062318617874881558261094899454408071345335844420144450891226007191374704820474341171395312812199095985963800532766775553

a = randint(2, p-2)
print(f"A --> B: p = {p}, g = {g}, ka = {pow(g, a, p)}")

b = randint(2, p-2)
print(f"B --> A: kb = {pow(g, b, p)}")

k = long_to_bytes(pow(g, a*b, p))

with open('text.png', 'rb') as file:
    text = file.read()

cipher = bytes([text[i] ^ k[i % len(k)] for i in range(len(text))])

with open('text.png.enc', 'wb') as file:
    file.write(cipher)
